import librosa
import numpy as np
import json

def analyze_audio_features(audio_path):
    """
    分析音频文件的特征，为谱面生成提供数据支持
    
    Args:
        audio_path (str): 音频文件路径
        
    Returns:
        dict: 包含音频特征的字典
    """
    try:
        # 加载音频文件
        y, sr = librosa.load(audio_path)
        
        # 获取音频时长
        duration = librosa.get_duration(y=y, sr=sr)
        
        # 提取节拍信息
        tempo, beats = librosa.beat.beat_track(y=y, sr=sr)
        
        # 估计主频率
        pitches, magnitudes = librosa.piptrack(y=y, sr=sr)
        
        # 计算频谱质心
        spectral_centroids = librosa.feature.spectral_centroid(y=y, sr=sr)[0]
        
        # 计算RMSE能量
        rmse = librosa.feature.rms(y=y)[0]
        
        # 计算色度特征
        chroma = librosa.feature.chroma_stft(y=y, sr=sr)
        
        # 按节拍采样特征
        beat_times = librosa.frames_to_time(beats, sr=sr)
        
        # 限制节拍数量以避免数据过载
        max_beats = 100  # 限制最多100个节拍
        if len(beat_times) > max_beats:
            # 采样节拍
            indices = np.linspace(0, len(beat_times)-1, max_beats, dtype=int)
            beat_times = beat_times[indices]
            beats = beats[indices]
        
        # 构建节拍特征
        beat_features = []
        for i, beat_time in enumerate(beat_times):
            # 获取节拍帧索引
            frame_idx = librosa.time_to_frames(beat_time, sr=sr)
            
            # 确保索引在范围内
            if frame_idx >= len(spectral_centroids) or frame_idx >= len(rmse):
                continue
                
            # 获取对应特征
            spectral_centroid = spectral_centroids[frame_idx] if frame_idx < len(spectral_centroids) else 0
            rms_energy = rmse[frame_idx] if frame_idx < len(rmse) else 0
            
            # 计算色度特征（平均值）
            if frame_idx < chroma.shape[1]:
                chroma_features = chroma[:, frame_idx].tolist()
            else:
                chroma_features = [0] * 12  # 12个半音
            
            # 计算主音高
            if frame_idx < pitches.shape[1]:
                # 获取该帧中幅度最大的频率
                max_magnitude_idx = np.argmax(magnitudes[:, frame_idx]) if magnitudes.shape[0] > 0 else 0
                pitch = pitches[max_magnitude_idx, frame_idx] if max_magnitude_idx < pitches.shape[0] else 0
            else:
                pitch = 0
            
            beat_features.append({
                'time': float(beat_time),
                'features': {
                    'spectral_centroid': float(spectral_centroid),
                    'rms_energy': float(rms_energy),
                    'pitch': float(pitch),
                    'chroma': chroma_features
                }
            })
        
        # 构建返回数据
        features = {
            'metadata': {
                'duration': float(duration),
                'tempo': float(tempo),
                'sample_rate': int(sr),
                'beat_count': len(beat_features)
            },
            'beats': beat_features
        }
        
        return features
        
    except Exception as e:
        print(f"音频分析出错: {e}")
        return None

def extract_features_for_chart_generation(audio_path):
    """
    为谱面生成提取专门优化的音频特征
    
    Args:
        audio_path (str): 音频文件路径
        
    Returns:
        dict: 优化后的音频特征，更适合谱面生成
    """
    # 获取基础特征
    features = analyze_audio_features(audio_path)
    
    if not features:
        return None
    
    # 进一步处理特征，使其更适合谱面生成
    optimized_features = {
        'metadata': features['metadata'].copy(),
        'beats': []
    }
    
    # 对每个节拍的特征进行优化处理
    for beat in features['beats']:
        optimized_beat = {
            'time': beat['time'],
            'features': {}
        }
        
        # 归一化能量值到0-1范围
        rms_energy = beat['features']['rms_energy']
        normalized_energy = np.clip(rms_energy * 2, 0, 1)  # 简单的归一化和放大
        optimized_beat['features']['rms_energy'] = float(normalized_energy)
        
        # 处理频谱质心（音高中心）
        spectral_centroid = beat['features']['spectral_centroid']
        # 将频谱质心映射到更易理解的范围
        optimized_beat['features']['spectral_centroid'] = float(spectral_centroid)
        
        # 简化色度特征为一个主导音高值
        chroma = beat['features']['chroma']
        if any(chroma):  # 如果有色度数据
            dominant_pitch_class = int(np.argmax(chroma))  # 获取最强烈的半音类
        else:
            dominant_pitch_class = 0
        optimized_beat['features']['dominant_pitch_class'] = dominant_pitch_class
        
        optimized_features['beats'].append(optimized_beat)
    
    return optimized_features

def save_features(features, output_file):
    """
    保存特征为JSON文件
    """
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(features, f, ensure_ascii=False, indent=2)

def main():
    """
    主函数
    """
    if len(sys.argv) < 2:
        print("使用方法: python simple_beat_analyzer.py <音频文件路径> [输出文件名]")
        print("示例: python simple_beat_analyzer.py track.mp3")
        print("示例: python simple_beat_analyzer.py track.mp3 features.json")
        return
    
    audio_file_path = sys.argv[1]
    
    if not os.path.exists(audio_file_path):
        print(f"错误: 找不到音频文件 '{audio_file_path}'")
        return
    
    if len(sys.argv) >= 3:
        output_file = sys.argv[2]
    else:
        output_file = "features.json"  # 简化输出文件名
    
    # 执行节拍级特征提取
    features = extract_simple_beat_features(audio_file_path)
    
    if features is None:
        print("特征提取失败")
        return
    
    # 保存结果
    save_features(features, output_file)
    
    print(f"\n简化节拍级音频特征提取完成!")
    print(f"特征数据已保存到: {output_file}")
    print(f"总共分析了 {features['metadata']['total_beats']} 个节拍点")
    print("\nJSON结构说明:")
    print("metadata:")
    print("  - audio_file: 音频文件路径")
    print("  - duration: 音频总时长(秒)")
    print("  - tempo: BPM节拍速度")
    print("  - total_beats: 总节拍数")
    print("  - beat_interval: 节拍间隔(秒)")
    print("beats: 节拍级特征数组，每个元素包含:")
    print("  - beat_index: 节拍索引")
    print("  - time: 节拍时间点(秒)")
    print("  - features: 节拍音频特征")
    print("    - rms_energy: 均方根能量(响度)")
    print("    - spectral_centroid: 频谱质心(音高中心)")
    print("    - spectral_bandwidth: 频谱带宽(音色宽度)")
    print("    - main_pitch_class: 主音高类别(0=C, 1=C#,..., 11=B)")
    print("    - pitch_strength: 主音高强度")
    print("    - chroma: 12个半音的强度值")
    print("  - context: 与前一个节拍的差异")
    print("    - rms_change: 响度变化")
    print("    - centroid_change: 音高中心变化")

if __name__ == "__main__":
    main()